# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import abc
import datetime
import functools
import logging
import signal
import typing

from recordclass import dataobject
from tornado.ioloop import PeriodicCallback

from coin.base.param_util import to_list
from coin.exchange.base.strategy_logger import StrategyLogProvider
from coin.strategy.accounting.pnl_balance_printer.factory import get_pnl_balance_printer_cls
from coin.strategy.accounting.pnl_balance_printer.util import (get_quote_currency, reserve_as_proto)
from coin.strategy.executor.util.cancel_all_executor import SimpleCancelAllExecutor
from coin.strategy.mm.base.bookmap import BookMap
from coin.strategy.mm.base.unary_multifeed_strategy_base import UnaryMultifeedStrategyBase
from coin.strategy.mm.subscription import FeedSubscriptionRequest
from coin.strategy.tool.fill_log_dumper import FillLogDumper


class StrategyDriverConfig(dataobject):
  currencies: typing.Any = []
  products: typing.Any = []
  sub_req: typing.Any = None
  og_params: typing.Any = None


class StrategyBase(metaclass=abc.ABCMeta):
  def __init__(self):
    self._strategy_logger = None

  @property
  def strategy_logger(self):
    return self._strategy_logger

  @strategy_logger.setter
  def strategy_logger(self, strategy_logger):
    assert self._strategy_logger is None, 'Strategy logger is already set.'
    self._strategy_logger = strategy_logger

  @abc.abstractmethod
  def prepare(self, feed_subsystem):
    raise NotImplementedError()

  @abc.abstractmethod
  def on_og_reset(self, order_gateway):
    raise NotImplementedError()

  @abc.abstractmethod
  def on_exit(self):
    raise NotImplementedError()

  @abc.abstractproperty
  def strategy_request(self):
    raise NotImplementedError()


class StrategyDriver(UnaryMultifeedStrategyBase):
  def __init__(self, strategy, driver_config, logger=None):
    self._logger = logger or logging.getLogger('StrategyDriver')
    self._strategy = strategy

    self._config = driver_config
    self._set_traits(self._config.sub_req)
    super().__init__(products=self._config.products,
                     currencies=self._config.currencies,
                     logger=self._logger)

    self._book = BookMap(self._products, logger=self._logger)

    reserve_map = self._gen_reserve_map()
    printer_cls = get_pnl_balance_printer_cls(self._config.sub_req)
    self._printer = printer_cls(self._products, self._book, reserve_map, logger=self._logger)

    self._fill_dumper = FillLogDumper(self._config.sub_req)

    self._order_canceled_after_og_ready = False
    self._exit_attempt = 0
    self._exit_executor = None

    self._strategy_logger = StrategyLogProvider(self._strategy.strategy_request)
    self._strategy_logger.write_request()
    self._strategy.strategy_logger = self._strategy_logger
    if reserve_map:
      self._strategy_logger.write_reserve(reserve_as_proto(self._config.sub_req, reserve_map))

  def register_signal_handler(self):
    signal.signal(signal.SIGTERM, self.on_signal)
    signal.signal(signal.SIGINT, self.on_signal)

  @property
  def strategy_logger(self):
    return self._strategy_logger

  @property
  def config(self):
    return self._config

  def _gen_reserve_map(self):
    reserve_map = {}
    '''
    TODO(jaewon):
    if self._config.trade_subreq.market_type == 'Spot':
      for profile in self._config.model_profile.values():
        currency = profile.trade_product.base.currency
        min_pos = profile.config_dict['min_pos']
        max_pos = profile.config_dict['max_pos']
        reserve = 0.5 * (min_pos + max_pos)
        resep[currency] = reserve
    '''
    return reserve_map

  def dump_pnl_balance(self):
    self._printer.print_pnl_balance()
    pnl_balances = self._printer.calculate_pnl_balance()
    if pnl_balances is not None:
      for pnl_balance in pnl_balances:
        self._strategy_logger.write_pnl_balance(pnl_balance=pnl_balance, quote=None, subreq_id=None)

  def prepare(self):
    assert self.feed_subsystem is not None
    self._strategy.prepare(self.feed_subsystem)

    for product in self._config.products:
      feed_sub_req = FeedSubscriptionRequest.create([product], self._config.sub_req)
      self.feed_subsystem.subscribe_book(feed_sub_req, functools.partial(self.on_book, product))

    # Start dumping pnl balance every 10 seconds
    PeriodicCallback(self.dump_pnl_balance, 10 * 1000).start()

  def on_og_reset(self):
    if not self._order_canceled_after_og_ready:
      self.order_gateway.cancel_all()
      self._logger.info('Orders canceled after OGs are ready.')
      self._order_canceled_after_og_ready = True

    self._printer.register_order_gateway(self.order_gateway)
    self._strategy.on_og_reset(order_gateway=self.order_gateway)
    self.order_gateway.register_fill_callback(self._fill_dumper.dump_fill)

    self._exit_executor = SimpleCancelAllExecutor()

  def on_book(self, product, book):
    self._book.set_book(product, book)

  def on_signal(self, sig, frame):
    self._logger.warning('Caught signal: %s', sig)
    self._ioloop.add_callback_from_signal(self.on_exit)

  def on_exit(self):
    self._strategy.on_exit()

    self._exit_attempt += 1
    if self._exit_attempt > 5:
      # Exit immediately.
      self._ioloop.add_callback(self._ioloop.stop)

    finished = True
    if self._exit_executor is not None:
      ogs = to_list(self.order_gateway)
      self._logger.info('Cancel all.')
      finished = self._exit_executor.cancel_all_on_exit(ogs)

    if finished:
      self._ioloop.add_callback(self._ioloop.stop)
    else:
      timeout = datetime.timedelta(seconds=2)
      self._ioloop.add_timeout(timeout, self.on_exit)
